from datasets import load_dataset
import glob, tiktoken
# 文件顶部
from tqdm import tqdm  # ✅ 从模块里取出同名函数

# ---------- 参数区 ----------
MODEL = "gpt-3.5-turbo"  # 也可 gpt-4 / gpt-4-turbo
MAX_CODE_TOKENS = 3000  # 控制代码长度，防止超长
MAX_OUTPUT = 1024  # 最大 diff token
TEMP = 0.2

# ----------------------------

encoding = tiktoken.encoding_for_model(MODEL)


def trim_code(code: str, max_tokens: int) -> str:
    """简单截断，保证不超 token"""
    tokens = encoding.encode(code)
    if len(tokens) <= max_tokens:
        return code
    return encoding.decode(tokens[:max_tokens])


def read_problem_file(instance):
    """把 patch 里第一个文件整段读出"""
    repo = instance["repo"].replace("/", "@")
    patch_path = instance["patch"].split(" a/")[1].split(" b/")[0] # 'astropy/modeling/separable.py'
    print(patch_path)
    candidates = glob.glob(f"repos/{repo}/**/{patch_path}", recursive=True)
    if not candidates:
        return ""
    with open(candidates[0], encoding="utf-8") as f:
        return f.read()


def build_prompt(problem: str, code: str) -> str:
    code = trim_code(code, MAX_CODE_TOKENS)
    return f"""Below is a GitHub issue and the corresponding source file.
            Issue:
            {problem}

            Source file:
            {code}
            Please generate a unified git diff that fixes the issue.
            Return only the diff block, starting with diff --git.
            """


# def call_openai(prompt: str) -> str:
#     rsp = client.chat.completions.create(
#         model=MODEL,
#         messages=[{"role": "user", "content": prompt}],
#         temperature=TEMPERATURE,
#         max_tokens=MAX_OUTPUT_TOKENS,
#     )
#     return rsp.choices[0].message.content or ""


def main():
    # 替换为你的本地文件路径
    ds = load_dataset(
        "parquet",
        data_files={  # 由于无法访问到Hugface，这里加载下载好的文件
            "test": r"D:\JunTuan\zxli26\0303_swebench\swe-bench-verified\data\test-00000-of-00001.parquet"
        }, split='test'
    )

    preds = {}

    for ins in tqdm(ds, desc="Generating patches"):
        instance_id = ins["instance_id"]
        print(instance_id)
        print("-----------------------------------------------------------------------")
        problem = ins["problem_statement"]
        repo = ins["repo"]
        code = read_problem_file(ins)
        prompt = build_prompt(problem, code)
        print("=======================================================================")
        print(prompt)
        break



if __name__ == '__main__':
    main()
